"""
/* Copyright 2018 The Enflame Tech Company. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
"""
# !/usr/bin/python
# coding=utf-8

import sys
if sys.version_info[0] < 3:
    import monotonic
    mono_time = monotonic.time.time
else:
    import time
    mono_time = time.monotonic
import json
import re
from collections import OrderedDict
try:
    import queue
except ImportError:
    import Queue as queue
import threading

from utils.flags import InfoDict
from utils.dtu_logger import LOGGER as logger

base_info = InfoDict().base_info


class RuntimeWatcher(object):
    """
        Collect time stamp throughout all the life cycle.
        It will also calculate performance index(fps, including average, min and max) before all end.
        Average value shows an objective performance index while min and max together show a fluctuating range.
    """
    def __init__(self, batch_size_per_step):
        self.batch_size_per_step = batch_size_per_step
        self.time_list = OrderedDict()
        self.session_time_list = list()
        self.duration_time_list = list()
        self.step_fps_list = list()
        self.start_time = queue.Queue()

    def watch(self, message=None):
        """
            Collect time stamp throughout all the life cycle.
            Current time will be recorded and stored once called.
            It will also calculate performance index(fps) before all end.
            Args:
                message. Message string should contain 'start' or 'end' as regular matching inside watcher
        """
        current_time = mono_time()
        if re.search("session", message):
            if re.search("start", message):
                self.start_time.put(current_time)
                self.session_time_list.append("{} session_start".format(int(round(current_time * 1e9))))
            if re.search("end", message):
                tmp_start = self.start_time.get()
                current_duration = current_time - tmp_start
                if current_duration > 0:
                    self.duration_time_list.append(current_duration)
                step_fps = self.batch_size_per_step / current_duration
                logger.info(
                    "{}:{}, session_start:{}, session_duration:{}, step_fps:{}".format(message,
                                                                                       current_time,
                                                                                       tmp_start,
                                                                                       current_duration,
                                                                                       step_fps))
                if step_fps > 0:
                    self.step_fps_list.append(step_fps)
                self.session_time_list.append("{} session_end".format(int(round(current_time * 1e9))))

        current_time = int(round(current_time * 1e9))
        self.time_list[message] = current_time

    def time_json(self):
        time_json = json.dumps(self.time_list, indent=4, sort_keys=False)
        return time_json

    def dump_json(self):
        """
            Save as a file in JSON format.
        """
        with open("benchmark_trace.log", "w") as f:
            for line in self.session_time_list:
                f.write(line + '\n')

    def calc_fps(self):
        """
           Calculate average fps.
           If only one step, return the value.
           If only two steps, return the second value.
           Otherwise, return average value of values excluding the first.
        """
        assert len(self.duration_time_list) > 0, 'Failed to obtain a full session cycle'
        if len(self.duration_time_list) == 1:
            logger.warn("One-step session run cannot calculate fps correctly")
            average_fps = self.step_fps_list[0]
        else:
            logger.info("Calculate fps excluding 1st step")
            total_duration = sum(self.duration_time_list[1:])
            average_duration = total_duration / (len(self.duration_time_list) - 1)
            average_fps = self.batch_size_per_step / average_duration
            average_fps = average_fps
        return average_fps

    def min_fps(self):
        """The min fps of all fps values. Equal to average and max when only one or two steps"""
        if len(self.step_fps_list) == 1:
            return self.step_fps_list[0]
        else:
            return min(self.step_fps_list[1:])

    def max_fps(self):
        """The max fps of all fps values. Equal to average and min when only one or two steps"""
        max_fps = max(self.step_fps_list)
        return max_fps

class EnflameThread(threading.Thread):
    def __init__(self):
        super(EnflameThread, self).__init__()

    def watcher(self, batch_size_per_step):
        """
            Creat a RuntimeWatcher object in another thread.
            Args:
                 batch_size_per_step: batch size per step, batch_per_cluster * 4 for 4c.
        """
        threading_watcher = RuntimeWatcher(batch_size_per_step)
        return threading_watcher


if __name__ == '__main__':
    import tensorflow as tf

    a = tf.random_normal([])
    b = tf.random_normal([])
    c = a + b
    init = tf.global_variables_initializer()
    watcher = EnflameThread().watcher(4)

    with tf.Session() as sess:
        watcher.watch("init start")
        sess.run(init)
        # watcher.watch("init end")
        for i in range(10):
            watcher.watch("step {} session start".format(i))
            sess.run(c)
            watcher.watch("step {} session end".format(i))
    print(watcher.time_json())
    watcher.dump_json()
